#include "core/frame.hpp"
#include "core/constants.hpp"

#include "yocto/code/utils.hpp"
#include "yocto/code/unroll.hpp"


namespace FastMD 
{
	Frame:: Frame( const size_t num_vectors, size_t num_vertices ) : 
	vectors( max_of(num_vectors,MIN_VECTORS) ),
	particles( max_of<size_t>(num_vertices,1) ),
	box(1,1,1),
	ibox(1,1,1),
	pbc_r( PBC_xyz ),
	pbc_f( PBC_xyz ),
	V(0),
	K(0),
	W(0),
	vector_( memory::kind<MemoryKind>::acquire_as<Vector>((size_t&)vectors) )
	{
		size_t iVec = 0;
		try 
		{
			while( iVec < vectors )
			{
				new ( &vector_[iVec] ) Vector( particles ); 
				++iVec;
			}
		}
		catch(...)
		{
			release(iVec);
			throw;
		}
	}
	
	Frame:: ~Frame() throw() { release( vectors ); }
	
	void Frame:: release( size_t n ) throw()
	{
		while( n > 0 )
		{
			vector_[--n].~Vector();
		}
		memory::kind<MemoryKind>::release_as<Vector>(vector_,(size_t&)vectors);
	}
	
	void Frame:: set_box( const V3D &b ) throw()
	{
		V3D & __box = (V3D&)box;
		V3D & _ibox = (V3D&)ibox;
		(_ibox = ( __box = b ) ).reciprocal();
	}
	
	
#define F2A(INDEX) a[INDEX].mul( r[INDEX].data )
	
	void Frame:: forces_to_accelerations() throw()
	{
		V3D       *a = &vector_[ ACCELERATIONS ][0];
		const V3D *r = &vector_[ POSITIONS     ][0];
		YOCTO_LOOP_FUNC_(particles,F2A,0);
	}
	
	void Frame:: update_K() throw()
	{
		const V3D *v = &vector_[ VELOCITIES ][0];
		K = 0;
#define KUPDATE(i) K += v[i].data * v[i].normsq()
		YOCTO_LOOP_FUNC_(particles,KUPDATE,0);
		K *= __HALF;
	}
	
	void Frame:: initialize_velocities( const Real temperature ) throw()
	{
		assert( temperature >= 0 );
		if( temperature > 0 )
		{
			//------------------------------------------------------------------
			//-- compute Kinetic Energy
			//------------------------------------------------------------------
			const Real   RT   =  R * temperature;
			const size_t dof  =  3*particles;
			K                 =  dof * RT * __HALF;
			
			//------------------------------------------------------------------
			//-- Create a Maxwell-Boltzmann distribution
			//------------------------------------------------------------------
			V3D  P;      //! linear momentum
			for( size_t i=0; i < particles; ++i )
			{
				const V3D &r   = vector_[ Frame::POSITIONS  ][i];
				V3D       &v   = vector_[ Frame::VELOCITIES ][i];
				const Real sig = __sqrt( RT * r.data ); //! RT/m
				v.x = sig * Normal();
				v.y = sig * Normal();
				v.z = sig * Normal();
				
				P.muladd( v.data, v );
				//std::cerr << "[" << v.x << " " << v.y << " " << v.z << "]" << std::endl;
			}
			
			//------------------------------------------------------------------
			//-- make a global zero linear momentum
			//-- an collect current kinetic energy
			//------------------------------------------------------------------
			const V3D dP( -__ONE/particles, P );
			Real      K1 = 0;
			for( size_t i=0; i < particles; ++i )
			{
				const V3D &r   = vector_[ Frame::POSITIONS  ][i];
				V3D       &v   = vector_[ Frame::VELOCITIES ][i];
				v.muladd( r.data, dP );
				K1 += v.data * v.normsq();
			}
			K1 *= __HALF;
			
			//------------------------------------------------------------------
			//-- rescale velocities
			//------------------------------------------------------------------
			const Real lam = __sqrt(K/K1);
			for( size_t i=0; i < particles; ++i )
			{
				V3D       &v   = vector_[ Frame::VELOCITIES ][i];
				v.mul(lam);
			}			
		}
		else
		{
			K = 0;
			for( size_t i=0; i < particles; ++i )
			{
				vector_[ Frame::VELOCITIES ][i].ldz();
			}
		}
		
		
	}
	
	void Frame:: rescale( const V3D &b ) throw()
	{
		const V3D fac( b.x * ibox.x, b.y * ibox.y, b.z * ibox.z );
		V3D & __box = (V3D&)box;
		V3D & _ibox = (V3D&)ibox;
		(_ibox = ( __box = b ) ).reciprocal();
		
		Vector &pos = vector_[ Frame::POSITIONS  ];
		for( size_t i=0; i < particles; ++i )
		{
			V3D &r = pos[i];
			r.x *= fac.x;
			r.y *= fac.y;
			r.z *= fac.z;
			pbc_r(r,box,ibox);
		}
	}
	
	void Frame:: thermal_friction()
	{
		const size_t n   = particles;
		Vector    &force = vector_[ ACCELERATIONS ];
		Vector    &v     = vector_[ VELOCITIES    ];
		
		//----------------------------------------------------------------------
		//compute the friction coefficient
		//----------------------------------------------------------------------
		Real zeta = 0;
		Real       num   = 0;
		Real       den   = 0;
		for( size_t i=0; i < n; ++i )
		{
			const V3D &v_i = v[i];
			num += v_i.dot( force[i] );
			den += v_i.normsq();
		}
		if( den > 0 )
		{
			zeta = -num / den;
			
			//------------------------------------------------------------------
			// propagate it
			//------------------------------------------------------------------
			for( size_t i=0; i < n; ++i )
			{
				force[i].muladd(zeta,v[i]);
			}
		}
	}
	
	void Frame:: thermal_update_K( const Real K_current ) throw()
	{
		const Real lam = __sqrt(K/K_current);
#define __UPDATE_V(INDEX) v[INDEX].mul(lam)
		Vector &v = vector_[ VELOCITIES ];
		YOCTO_LOOP_FUNC_(particles,__UPDATE_V,0);
		
	}
	
}

